#!/usr/bin/env python

__version__ = "$Id$"

import logging
from contacts import *
import contacts
import unittest
import doctest
import sqlite3

TEST_CONTACT = """HBPLUS Hydrogen Bond Calculator v 3.15            Nov 13 17:34:59 CET 2009
(c) I McDonald, D Naylor, D Jones and J Thornton 1993 All Rights Reserved.
Citing HBPLUS in publications that use these results is condition of use.
9TTT <- Brookhaven Code "9TTT.pdb.new" <- PDB file                         
<---DONOR---> <-ACCEPTOR-->    atom                        ^               
c    i                          cat <-CA-CA->   ^        H-A-AA   ^      H- 
h    n   atom  resd res      DA  || num        DHA   H-A  angle D-A-AA Bond
n    s   type  num  typ     dist DA aas  dist angle  dist       angle   num
B0303-HOH O   B0001-  A O5* 3.17 HH  -2 -1.00  -1.0 -1.00  -1.0 126.2     1
B0285-HOH O   B0001-HOH N1  2.70 HH  -2 -1.00  -1.0 -1.00  -1.0 102.5     2
A0170-ARG NH2 B0002-  G O2P 2.55 SH  -2 -1.00 112.7  1.99 106.6 117.4     3
A0265-HOH O   B0002-  G O2P 2.98 HH  -2 -1.00  -1.0 -1.00  -1.0 127.5     4
A0279-HOH O   B0002-  U O2P 2.79 HH  -2 -1.00  -1.0 -1.00  -1.0 118.7     5
A0180-ARG NH2 B0002-LYS N7  2.71 SH  -2 -1.00 155.7  1.77 129.4 123.3     6
C0061-  G N4  B0002-  G O6  2.90 HH  -2 -1.00 169.9  1.91 122.7 122.0     7
A0180-ARG NH1 B0002-  G O6  2.91 SH  -2 -1.00 169.0  1.93 129.3 132.0     8
B0002-  G N1  C0061-  C N3  2.84 HH  -2 -1.00 178.0  1.84 120.7 117.4     9
B0002-  G N2  C0061-  U O2  2.67 HH  -2 -1.00 171.4  1.68 119.5 122.4    10
B0002-ARG N2  C0061-ARG O2  2.67 HH  -2 -1.00 171.4  1.68 119.5 122.4    11
"""
class ContactFileTest(unittest.TestCase):
    def setUp(self, *args):
    	#unittest.TestCase.__init__(self, *args)
    	
    	DB = "testdb.sqlite"
    	try:
    		os.remove(DB)
    	except OSError: # no such file
    		pass
    	self.c = [ContactFile("../pdb/3DVO.pdb.hb2", bond_type="H", DB_NAME = DB)]
    	#print "3DVO"
        self.c.append(ContactFile("../pdb/2B0E.pdb.hb2", bond_type="H", DB_NAME = DB))
        #print "2B0E"
        self.c.append(ContactFile("../pdb/2B0E.pdb.nb2", bond_type="N",DB_NAME = DB))
        #print "2B0E N"
        self.c.append(ContactFile("../pdb/1AAY.pdb.hb2", bond_type="H", DB_NAME = DB))
        #print "1AAY"
        self.c.append(ContactFile("../pdb/1AAY.pdb.nb2", bond_type="N", DB_NAME = DB))
        #print "1AAY N"
        testfile = open("TEST.hb2", "w")
        testfile.write(TEST_CONTACT)
        testfile.close()
        self.c.append(ContactFile("TEST.hb2", bond_type="H", DB_NAME = DB))

        
        for con in self.c:
        	#print "    %s %s    " % (con.pdb, con.type)
        	con.save()
        #self.conn.row_factory = sqlite3.Row
        self.conn = sqlite3.connect("testdb.sqlite")
        self.cur = self.conn.cursor()
    
    def testInit(self):
        self.failUnless(self.c[1].type == "H")
        self.failUnless(self.c[2].type == "N")
        self.failUnless(self.c[3].type == "H")
        self.failUnless(self.c[4].type == "N")
        self.failUnless(self.c[1].pdb == "2B0E")
        self.failUnless(self.c[2].pdb == "2B0E")
        self.failUnless(self.c[3].pdb == "1AAY")
        self.failUnless(self.c[4].pdb == "1AAY")
        self.failUnless(self.c[5].pdb == "9TTT")
    
    def testExclusion(self):
        self.failUnless(len(self.c[5]) == 5)
        ids = set((1, 3, 4, 5, 8))
        not_excluded = set(x.id for x in self.c[5])
        self.failUnless(ids == not_excluded)
    
    def testSave(self):
        SQL =			"""SELECT * FROM
                         contacts_9TTT AS contacts
                            JOIN atoms_9TTT AS donor
                                ON donor.id = contacts.donor_id
                            JOIN atoms_9TTT AS acceptor
                                ON acceptor.id = contacts.acceptor_id"""
        model_reply = [(1, u'H', 1, 2, 3.1699999999999999, u'H', u'H', -2, -1, -1, -1, -1, 126.2, u'D', 1, u'B', 303, u'-', u'HOH', u'O', 2, u'B', 1, u'-', u'A', u'O5*'),
                            (2, u'H', 3, 4, 2.5499999999999998, u'S', u'H', -2, -1, 112.7, 1.99, 106.59999999999999, 117.40000000000001, u'D', 3, u'A', 170, u'-', u'ARG', u'NH2', 4, u'B', 2, u'-', u'G', u'O2P'),
                            (3, u'H', 5, 4, 2.98, u'H', u'H', -2, -1, -1, -1, -1, 127.5, u'D', 5, u'A', 265, u'-', u'HOH', u'O', 4, u'B', 2, u'-', u'G', u'O2P'),
                            (4, u'H', 6, 7, 2.79, u'H', u'H', -2, -1, -1, -1, -1, 118.7, u'D', 6, u'A', 279, u'-', u'HOH', u'O', 7, u'B', 2, u'-', u'U', u'O2P'),
                            (5, u'H', 8, 9, 2.9100000000000001, u'S', u'H', -2, -1, 169, 1.9299999999999999, 129.30000000000001, 132, u'D', 8, u'A', 180, u'-', u'ARG', u'NH1', 9, u'B', 2, u'-', u'G', u'O6')]
        # print [x[0] for x in self.cur.description]
        self.db_query(SQL, model_reply)
        
        SQL = "select type, count(*) from contacts_1AAY GROUP BY type ORDER BY type"
        model_reply = [(u'H',250), (u'N',965)]
        self.db_query(SQL, model_reply)
        
        
        
    def tearDown(self):
        self.conn.close()
        os.remove("testdb.sqlite")
        os.remove("TEST.hb2")
        
    def db_query(self, SQL, model_reply):
        self.cur.execute(SQL)
        for r in model_reply:
            self.failUnlessEqual(self.cur.fetchone(), r)
		
if __name__ == "__main__":
    doctest.testmod(contacts)
    #print Contact(('A', 390, '-', 'HOH', ' O  ', 'B', 1, '-', '  A', ' O5*', 3.2000000000000002, 'H', 'H', -2, -1.0, -1.0, -1.0, -1.0, 135.19999999999999, 1), "TEST")

    unittest.main()
